package com.algorithm.trie;

import java.util.TreeMap;

/**
 * @Author: DOWN
 * @Date: 2022/5/7 15:47
 * @Comment: 纯英文字典树
 */
public class TrieTree {
    /**
     * 内部 字母节点类
     */
    private class Node {
        /**
         * 是否为单词结尾
         */
        public boolean isWord;
        /**
         * 下一个节点
         */
        public TreeMap<Character, Node> next;

        /**
         * 初始化节点
         *
         * @param isWord 是否为单词结尾
         */
        public Node(boolean isWord) {
            this.isWord = isWord;
            next = new TreeMap<>();
        }

        /**
         * 默认不是单词结尾
         */
        public Node() {
            this(false);
        }
    }

    /**
     * 根节点
     */
    private Node root;
    /**
     * 长度
     */
    private int size;

    /**
     * 初始化字典树
     */
    public TrieTree() {
        root = new Node();
        size = 0;
    }

    /**
     * 获取字典树长度
     *
     * @return int
     */
    public int getSize() {
        return this.size;
    }

    /**
     * 添加一个新的单词
     *
     * @param word 英文单词
     */
    public void add(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            //如果不存在，则新增节点
            if (cur.next.get(c) == null) {
                cur.next.put(c, new Node());
            }
            cur = cur.next.get(c);
        }
        //判断当前是否为单词结尾
        if (!cur.isWord) {
            cur.isWord = true;
            //是，则长度加1
            size++;
        }
    }

    /**
     * 添加一个新的单词 - 递归
     *
     * @param word 英文单词
     */
    public void ReAdd(String word) {
        ReAdd(0, word, root);
    }

    /**
     * 添加一个新的单词 - 递归方法
     *
     * @param index 单词索引
     * @param word  单词
     * @param node  节点
     */
    public void ReAdd(int index, String word, Node node) {
        if (index >= word.length()) {
            return;
        }
        Node cur = node;
        char c = word.charAt(index);
        //如果不存在，则新增节点
        if (cur.next.get(c) == null) {
            cur.next.put(c, new Node());
            cur = cur.next.get(c);
            //判断当前是否为单词结尾
            if (!cur.isWord) {
                cur.isWord = true;
                //是，则长度加1
                size++;
            }
        } else {
            ReAdd(++index, word, cur.next.get(c));
        }
    }

    /**
     * 查询单词
     *
     * @param word 单词
     * @return boolean
     */
    public boolean contains(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (cur.next.get(c) == null) {
                return false;
            }
            cur = cur.next.get(c);
        }
        return cur.isWord;
    }
    /**
     * 查询前缀，每个单词也是本身的前缀。
     *
     * @param prefix 前缀
     * @return boolean
     */
    public boolean isPreFix(String prefix) {
        Node cur = root;
        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            if (cur.next.get(c) == null) {
                return false;
            }
            cur = cur.next.get(c);
        }
        //单词也是本身的前缀
        return true;
    }
}
